package com.wzs.rpc.registry;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.cron.CronUtil;
import cn.hutool.cron.task.Task;
import cn.hutool.json.JSONUtil;
import com.wzs.rpc.config.RegistryConfig;
import com.wzs.rpc.model.ServiceMetaInfo;
import io.etcd.jetcd.*;
import io.etcd.jetcd.kv.GetResponse;
import io.etcd.jetcd.kv.PutResponse;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

/**
 * 实现etcd注册中心，java客户端使用方式参考：<a href="https://www.cnblogs.com/wuyongyin/p/18789780">...</a>
 */
@Slf4j
public class EtcdRegistry implements Registry{
    private Client client;
    private static final String ETCD_ROOT_PATH = "/rpc/";
    //本机注册的节点key集合（用于维护续期）
    private Set<String> localRegisterNodeKeySet = new HashSet<>();
    //保存已监听的key，防止同一个key被重复监听
    private Set<String> watchingKeySet = new HashSet<>();
    //缓存服务列表(key为serviceKey，即ServiceMetaInfo.getServiceKey())
    private Map<String, List<ServiceMetaInfo>> serviceCache = new ConcurrentHashMap<>();

    //应该区分一下是provider还是consumer，如果是后者就没必要开启心跳检测
    @Override
    public void init(RegistryConfig registryConfig) {
        log.info("当前节点上线，节点地址为：{}", registryConfig.getAddress());
        client = Client.builder()
                .endpoints(registryConfig.getAddress())
                .connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
                .build();
    }

    @Override
    public void init(RegistryConfig registryConfig, boolean heartBeat) {
        init(registryConfig);
        if (heartBeat) {
            //开启心跳检测
            heartBeat();
        }
    }

    @Override
    public void heartBeat() {
        CronUtil.schedule("*/10 * * * * *", new Task() {
            @Override
            public void execute() {
                for (String registerKey : localRegisterNodeKeySet) {
                    try {
                        List<KeyValue> keyValues = client.getKVClient()
                                .get(ByteSequence.from(registerKey, StandardCharsets.UTF_8))
                                .get()
                                .getKvs();
                        //说明该节点已过期
                        if (CollUtil.isEmpty(keyValues)) {
                            continue;
                        }
                        //节点未过期，续签
                        KeyValue keyValue = keyValues.get(0); //因为是精准匹配，所以只有一个结果
                        String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                        ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
                        register(serviceMetaInfo);
                    } catch (Exception e) {
                        log.error("{} 节点续期失败", registerKey, e);
                        throw new RuntimeException("节点续期失败", e);
                    }
                }
            }
        });

        //支持秒级定时任务
        CronUtil.setMatchSecond(true);
        CronUtil.start();
    }

    //这里就会发现，我们最好能把关于该服务的完整信息存进去，而不仅仅是节点地址
    //比如还包括服务的版本号、服务分组等信息，所以最好将有关服务的信息封装成一个DO
    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws Exception{
        log.debug("注册服务:{}", serviceMetaInfo);
        Lease leaseClient = client.getLeaseClient();
        // 创建一个 300 秒的租约
        long leaseId = leaseClient.grant(300).get().getID();

        // 设置要存储的键值对
        String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence key = ByteSequence.from(registerKey, StandardCharsets.UTF_8);
        ByteSequence value = ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);

        // 将键值对与租约关联起来
        PutOption putOption = PutOption.builder().withLeaseId(leaseId).build();
        client.getKVClient().put(key,value,putOption).get();

        //添加节点信息到key集合
        localRegisterNodeKeySet.add(registerKey);
    }

    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) throws Exception{
        log.info("注销服务:{}", serviceMetaInfo);
        String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence key = ByteSequence.from(registerKey, StandardCharsets.UTF_8);
        client.getKVClient().delete(key).get();
        
        //从key集合中删除节点信息
        localRegisterNodeKeySet.remove(registerKey);
    }

    //serviceKey=serviceName:serviceVersion
    @Override
    public List<ServiceMetaInfo> getService(String serviceKey) {
        //优先从serviceCache中获取服务
        if (serviceCache.containsKey(serviceKey)) {
            return serviceCache.get(serviceKey);
        }
        // 前缀搜索，结尾一定要加 '/'
        String searchPrefix = ETCD_ROOT_PATH + serviceKey + "/";
        log.debug("获取服务列表: {}", searchPrefix);
        try {
            GetOption getOption = GetOption.builder().isPrefix(true).build();
            List<KeyValue> keyValues = client.getKVClient()
                    .get(ByteSequence.from(searchPrefix, StandardCharsets.UTF_8), getOption)
                    .get()
                    .getKvs();

            List<ServiceMetaInfo> serviceMetaInfos = keyValues.stream()
                    .map(keyValue -> {
                        String key = keyValue.getValue().toString(StandardCharsets.UTF_8);
                        String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                        ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
                        watch(key, serviceMetaInfo);
                        return serviceMetaInfo;
                    })
                    .collect(Collectors.toList());
            serviceCache.put(serviceKey,serviceMetaInfos);
            return serviceMetaInfos;
        } catch (Exception e) {
            log.error("获取服务列表失败: {}", searchPrefix);
            throw new RuntimeException("获取服务列表失败", e);
        }
    }

    @Override
    public void destroy() {
        log.info("当前节点下线");
        for (String registerKey : localRegisterNodeKeySet) {
            try {
                client.getKVClient()
                        .delete(ByteSequence.from(registerKey, StandardCharsets.UTF_8))
                        .get();
            } catch (Exception e) {
                log.error(registerKey + "节点下线失败", e);
                throw new RuntimeException(registerKey + "节点下线失败", e);
            }
        }
        if (client!=null) {
            client.close();
        }
    }

    @Override
    public void watch(String registerKey, ServiceMetaInfo serviceMetaInfo) {
        if (watchingKeySet.contains(registerKey)) {
            return;
        }
        watchingKeySet.add(registerKey);
        String serviceKey = serviceMetaInfo.getServiceKey();
        client.getWatchClient()
                .watch(ByteSequence.from(registerKey, StandardCharsets.UTF_8), reponse -> {
                    for (WatchEvent event : reponse.getEvents()) {
                        switch (event.getEventType()) {
                            //todo：可能会出现NP异常
                            case DELETE:
                                serviceCache.get(serviceKey).remove(serviceMetaInfo);
                                break;
                            case PUT:
                                serviceCache.get(serviceKey).add(serviceMetaInfo);
                                break;
                        }
                    }
                });
    }
}
